#define _GNU_SOURCE
#include <string.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <sys/un.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>
#include <fcntl.h>
#include <glib.h>
#include <math.h>
#include <glib/gstdio.h>
#include <libvirt/libvirt.h>
#include <libvirt/virterror.h>
#include <libvirt/libvirt-qemu.h>
#include <json-c/json.h>

#include "private.h"
#include "driver/driver_wrapper.h"
#include "driver/memory_cache.h"
#include "driver/kvm/kvm.h"
#include "driver/kvm/kvm_private.h"

/* Memory Layout
ASP侧：
    0000000080000000-00000000ffffffff (prio 0, ram): riscv_virt_board.ram
    0000000100000000-000000017fffffff (prio 0, ram): {0c39ee18-6ac1-4ad4-96ea-cd4603ce51b8}-ASP-0x100000000-Host-0x80000000
Host侧：
    0000000080000000-00000000ffffffff (prio 0, ram): {0c39ee18-6ac1-4ad4-96ea-cd4603ce51b8}-ASP-0x100000000-Host-0x80000000
 */
#define CPU_PHY_MEM_STRART      0x80000000
#define CPU_PHY_MEM_END         0x100000000
#define CPU_PHY_MEM_SIZE        (CPU_PHY_MEM_END - CPU_PHY_MEM_STRART)

#define SEE_CPU_PHY_MEM_START   0x100000000
#define SEE_CPU_PHY_MEM_OFFSET  (SEE_CPU_PHY_MEM_START - CPU_PHY_MEM_STRART)

#define PAGE_SIZE               (1 << 12)
#define PAGE_ALIGN(x)           (x & ~(PAGE_SIZE - 1))
#define PAGE_OFFSET(x)          (x & (PAGE_SIZE - 1))

#define SOCKET_PATH             "/tmp/vmi.sock"

/* same with vmi-daemon.c */
#define VMI_MESSAGE_MAX_ARGS    6
struct vmi_message {
    int msg_type;
    int msg_argc;
    uint64_t msg_args[VMI_MESSAGE_MAX_ARGS];
};

enum vmi_message_type {
    VMI_MSG_GET_MEM_RANGE = 0,
    VMI_MSG_GET_GPA_DATA,
    VMI_MSG_ERR,
};

#ifdef HAVE_LIBVMI_REQUEST
#  include <qemu/libvmi_request.h>
#else

// request struct matches a definition in qemu source code
struct request {
    uint64_t type;   // 0 quit, 1 read, 2 write, ... rest reserved
    uint64_t address;   // address to read from OR write to
    uint64_t length;    // number of bytes to read OR write
};

#endif // !HAVE_LIBVMI_REQUEST

enum segment_type {
    SEGMENT_SELECTOR,
    SEGMENT_BASE,
    SEGMENT_LIMIT,
    SEGMENT_ATTR
};

//----------------------------------------------------------------------------
// Helper functions

static int init_socket(vmi_instance_t vmi) {
    kvm_instance_t *kvm = kvm_get_instance(vmi);
    if (kvm->socket_fd != -1) // socket is initialized
        return 0;

    errprint(">>> luowu. [%s] for VM. initialized to connect VMI Daemon\n", __func__);
    kvm->socket_fd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (kvm->socket_fd == -1) {
        errprint(">>> luowu. [%s] Socket creation failed", __func__);
        return -1;
    }

    struct sockaddr_un server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sun_family = AF_UNIX;
    strncpy(server_addr.sun_path, SOCKET_PATH, sizeof(server_addr.sun_path) - 1);

    if (connect(kvm->socket_fd, (struct sockaddr*)&server_addr, sizeof(struct sockaddr_un)) == -1) {
        errprint(">>> luowu. [%s] Connection failed", __func__);
        close(kvm->socket_fd);
        kvm->socket_fd = -1;
        return -1;
    }

    errprint(">>> luowu. [%s] succeed\n", __func__);
    return 0;
}

static int see_get_cpu_phy_mem(addr_t paddr_in_see, uint32_t length, char *out_buf)
{
    unsigned long *map_base;
    addr_t paddr_mapped = 0;
    uint32_t mapped_length = 0, paddr_offset = 0;
    int fd = open("/dev/mem", O_RDWR | O_SYNC);
    if (fd == -1) {
        errprint(">>> luowu. [%s] open /dev/mem failed\n", __func__);
        return -1;
    }

    paddr_mapped = PAGE_ALIGN(paddr_in_see);
    paddr_offset = PAGE_OFFSET(paddr_in_see);

    mapped_length = length + paddr_offset;
    if (PAGE_OFFSET(mapped_length) != 0)
        mapped_length += PAGE_SIZE;

    map_base = mmap(NULL, mapped_length, PROT_READ | PROT_WRITE , MAP_SHARED, fd, paddr_mapped);
    if (map_base == NULL) {
        errprint(">>> luowu. [%s] mmap failed\n", __func__);
        close(fd);
        return -1;
    }

    memcpy(out_buf, (char*)map_base + paddr_offset, length);
    munmap(map_base, mapped_length);

    close(fd);
    return 0;
}

static int see_request_vmi_daemon(vmi_instance_t vmi, char *msg, uint32_t msg_len, char *out_buf, uint32_t *out_len) {
    if (init_socket(vmi) != 0)
        return -1;

    write(kvm_get_instance(vmi)->socket_fd, msg, msg_len);

    /* wait for reply */
    ssize_t num_bytes = read(kvm_get_instance(vmi)->socket_fd, out_buf, *out_len);
    if (num_bytes == -1) {
        errprint(">>> luowu. [%s] read failed\n", __func__);
        return -1;
    }
    if (num_bytes < (ssize_t)sizeof(struct vmi_message)) {
        errprint(">>> luowu. [%s] bad msg from VMI Daemon\n", __func__);
        return -1;
    }
    if (((struct vmi_message*)out_buf)->msg_type == VMI_MSG_ERR) {
        errprint(">>> luowu. [%s] error from VMI Daemon\n", __func__);
        return -1;
    }

    errprint(">>> luowu. [%s] for VM. data replied by VMI Daemon is: 0x%lx (show 8 bytes)\n", __func__,
        num_bytes > 8 ? *(uint64_t*)out_buf : 0);

    *out_len = num_bytes;
    return 0;
}

void *
kvm_get_memory(
    vmi_instance_t vmi,
    addr_t paddr,
    uint32_t length)
{
    int id = kvm_get_id(vmi);
    if (id != HOST_ID && id != VM_ID) {
         errprint(">>> luowu. [%s] bad domainid %d\n", __func__, id);
         return NULL;
    }

    errprint(">>> luowu. [%s] paddr=0x%lx length=0x%x\n", __func__, paddr, length);
    char *buf = g_try_malloc0(length + 1);
    if (!buf)
        return NULL;

    if (paddr >= CPU_PHY_MEM_END || paddr < CPU_PHY_MEM_STRART) {
        errprint(">>> luowu. [%s] bad paddr\n", __func__);
        g_free(buf);
        return NULL;
    }

    if (id == HOST_ID) {
        if (see_get_cpu_phy_mem(paddr + SEE_CPU_PHY_MEM_OFFSET, length, buf) != 0) {
            errprint(">>> luowu. [%s] see_get_cpu_phy_mem failed\n", __func__);
            g_free(buf);
            return NULL;
        }
    } else if (id == VM_ID) {
        struct vmi_message msg;
        memset(&msg, 0, sizeof(msg));

        msg.msg_type = VMI_MSG_GET_GPA_DATA;
        msg.msg_argc = 2;
        msg.msg_args[0] = paddr;
        msg.msg_args[1] = length;

        uint32_t out_len = sizeof(msg) + length;
        char *out_buf = (char*)malloc(out_len); // the data content is attached following msg
        if (out_buf == NULL) {
            g_free(buf);
            return NULL;
        }
        memset(out_buf, 0, out_len);

        errprint(">>> luowu. [%s] for VM. request VMI Daemon to obtain data in GPA 0x%lx...\n", __func__, paddr);
        if (see_request_vmi_daemon(vmi, (char*)&msg, sizeof(msg), (char*)out_buf, &out_len) != 0 ||
            out_len != sizeof(msg) + length) {
            g_free(buf);
            free(out_buf);
            return NULL;
        }
        struct vmi_message *reply = (struct vmi_message*)out_buf;
        if (reply->msg_type != VMI_MSG_GET_GPA_DATA || reply->msg_argc != 2 ||
            reply->msg_args[0] != paddr || reply->msg_args[1] != length) {
            g_free(buf);
            free(out_buf);
            return NULL;
        }
        memcpy(buf, (char*)out_buf + sizeof(struct vmi_message), length);
        free(out_buf);
        
        errprint(">>> luowu. [%s] for VM. VMI Daemon tells us the data in GPA 0x%lx is: 0x%lx (show 8 bytes)\n", __func__,
            paddr, length > 4 ? *(uint64_t*)buf : 0);
        return buf;
    }
    return buf;
}

void
kvm_release_memory(
    vmi_instance_t UNUSED(vmi),
    void *memory,
    size_t UNUSED(length))
{
    if (memory)
        free(memory);
}

status_t
kvm_put_memory(vmi_instance_t vmi,
               addr_t paddr,
               uint32_t length,
               void *buf)
{
    errprint(">>> luowu. [%s] \n", __func__);
    return VMI_FAILURE;
}

status_t
kvm_init(
    vmi_instance_t vmi,
    uint32_t UNUSED(init_flags),
    vmi_init_data_t* UNUSED(init_data))
{
    errprint(">>> luowu. [%s] \n", __func__);
    kvm_instance_t *kvm = g_try_malloc0(sizeof(kvm_instance_t));
    if (!kvm)
        return VMI_FAILURE;

    // still create wrapper to avoid null ptr exception. will be removed finally
    if ( VMI_FAILURE == create_libvirt_wrapper(kvm) ) {
        g_free(kvm);
        return VMI_FAILURE;
    }

    // cannot connect to libvirt
    kvm->conn = NULL;

    vmi->driver.driver_data = (void*)kvm;

    return VMI_SUCCESS;
}

status_t
kvm_init_vmi(
    vmi_instance_t vmi,
    uint32_t init_flags,
    vmi_init_data_t* init_data)
{
    errprint(">>> luowu. [%s] \n", __func__);
    kvm_instance_t *kvm = kvm_get_instance(vmi);

    kvm->dom = NULL; // does not have it
    kvm->socket_fd = -1;
    vmi->vm_type = NORMAL;

    vmi->num_vcpus = 2;

    memory_cache_destroy(vmi);
    memory_cache_init(vmi, kvm_get_memory, kvm_release_memory, 1);

    (void)init_flags;
    (void)init_data;

    return VMI_SUCCESS;
}

void
kvm_destroy(
    vmi_instance_t vmi)
{
    errprint(">>> luowu. [%s] \n", __func__);
    kvm_instance_t *kvm = kvm_get_instance(vmi);
    if (!kvm)
        return;

    if (kvm->dom)
        kvm->libvirt.virDomainFree(kvm->dom);

    if (kvm->conn)
        kvm->libvirt.virConnectClose(kvm->conn);

    if (kvm->socket_fd != -1)
        close(kvm->socket_fd);

    dlclose(kvm->libvirt.handle);
    g_free(kvm);
}

status_t
kvm_get_memsize(
    vmi_instance_t vmi,
    uint64_t *allocated_ram_size,
    addr_t *maximum_physical_address)
{
#ifdef ENABLE_SAFETY_CHECKS
    if (!allocated_ram_size || !maximum_physical_address)
        return VMI_FAILURE;
#endif
    int id = kvm_get_id(vmi);
    if (id == HOST_ID) {
        *allocated_ram_size = CPU_PHY_MEM_SIZE; // bytes
        *maximum_physical_address = CPU_PHY_MEM_END;
        return VMI_SUCCESS;
    } else if (id == VM_ID) {
        struct vmi_message msg, reply;
        memset(&msg, 0, sizeof(msg));
        memset(&reply, 0, sizeof(reply));

        msg.msg_type = VMI_MSG_GET_MEM_RANGE;
        uint32_t out_len = sizeof(reply);
        if (see_request_vmi_daemon(vmi, (char*)&msg, sizeof(msg), (char*)&reply, &out_len) != 0 ||
            out_len != sizeof(reply) || reply.msg_type != VMI_MSG_GET_MEM_RANGE || reply.msg_argc != 2)
            return VMI_FAILURE;
        
        *allocated_ram_size = reply.msg_args[1] - reply.msg_args[0]; // bytes
        *maximum_physical_address = reply.msg_args[1];
        errprint(">>> luowu. [%s] for VM. VMI Daemon tells us the GPA range is: 0x%lx-0x%lx\n", __func__,
            reply.msg_args[0], reply.msg_args[1]);
        return VMI_SUCCESS;
    }

    return VMI_FAILURE;
}

bool kvm_check_pt_gfn_range(vmi_instance_t vmi, addr_t gfn)
{
    errprint(">>> luowu. [%s] \n", __func__);
    int id = kvm_get_id(vmi);
    if (id == HOST_ID) {
        if (gfn < (CPU_PHY_MEM_STRART >> 12) || gfn >= (CPU_PHY_MEM_END >> 12))
            return false;
        return true;
    } else if (id == VM_ID) {
        errprint(">>> luowu. [%s] for VM. Not Implemented\n", __func__);
        return false;
    }
    return false;
}

status_t
kvm_get_vcpureg(
    vmi_instance_t vmi,
    uint64_t *value,
    reg_t reg,
    unsigned long vcpu)
{
    errprint(">>> luowu. [%s]. Not Implemented\n", __func__);
    return VMI_FAILURE;
}

status_t
kvm_pause_vm(
    vmi_instance_t vmi)
{
    errprint(">>> luowu. [%s] Not Implemented. Directly return true\n", __func__);
    return VMI_SUCCESS;
}

status_t
kvm_resume_vm(
    vmi_instance_t vmi)
{
    errprint(">>> luowu. [%s] Not Implemented. Directly return true\n", __func__);
    return VMI_SUCCESS;
}
